// author: Shelman
// description: 封装的使用posix接口实现的TCP客户端

#include "TcpClient.h"
#include <iostream>
#include <cstring>
#include <chrono>

#ifdef _WIN32
#include <winsock2.h>
#include <ws2tcpip.h>
#include <windows.h>
#pragma comment(lib, "ws2_32.lib")
#else
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <unistd.h>
#include <fcntl.h>
#include <errno.h>
#include <poll.h>
#endif

#include "Common/CommLogDef.h"

#ifdef _WIN32
#define SOCK_ERR        SOCKET_ERROR
#else
#define SOCK_ERR        (-1)
#endif


/*****************************************************
 * 功能：构造函数
 * ***************************************************/
TcpClient::TcpClient()
    : m_bConnected(false)
    , m_bRunning(false)
{
#ifdef _WIN32
    // 初始化Winsock
    WSADATA wsaData;
    WSAStartup(MAKEWORD(2, 2), &wsaData);
#endif
}

/*****************************************************
 * 功能：析构函数
 * ***************************************************/
TcpClient::~TcpClient()
{
    disconnect();

#ifdef _WIN32
    // 清理Winsock
    WSACleanup();
#endif
}

/*****************************************************
 * 功能：不主动调用read接口获取数据时，设置接收数据的回调，
 * 输入参数：
 *      pfnCallback：数据回调函数
 * ***************************************************/
void TcpClient::setDataRecvCallback(const DataReceivedCallback &pfnCallback)
{
    std::lock_guard<std::mutex> lock(m_mt4PfnCallback);
    m_pfn4RecvDataCb = pfnCallback;
    return;
}

/*****************************************************
 * 功能：与服务端进行连接
 * 输入参数：
 *      stSrvConnInfo：服务端的连接信息
 *      bAutoReconnect：是否自动重连，true -- 自动重连，false -- 不自动重连
 * 返回值：
 *      bool：连接结果，true -- 连接成功，false -- 连接失败
 * ***************************************************/
bool TcpClient::connect(const NetworkConnInfo &stSrvConnInfo, bool bAutoReconnect)
{
    if (!m_bConnected) {
        // 设置数据标志等
        m_stSrvConnInfo = stSrvConnInfo;
        m_bAutoReconn = bAutoReconnect;

        // 初始化Socket
        initSocket();

        // 连接服务端
        if (!connectToServer()) {
            if (!m_bAutoReconn) {
                // 不需要自动重连时，直接返回连接失败
                COMMONLOG_ERROR("%s", "Connection timeout or error.");
                cleanup();
                return false;
            }
        }

        // 设置了自动重连，失败时，可在异步线程中一直尝试重连
        m_bRunning = true; // 设置标志
        // 启动异步线程
        m_th4AsyncIO = std::thread(&TcpClient::onAsyncIOThread, this);

        return true;
    } else {
        COMMONLOG_ERROR("%s", "Please diconnect first!");
    }

    return false;
}

/*****************************************************
 * 功能：断开与服务端的连接
 * ***************************************************/
void TcpClient::disconnect()
{
    // 清理线程资源
    if (m_bRunning) {
        m_bRunning = false;
        // 等待线程结束
        if (m_th4AsyncIO.joinable()) {
            m_th4AsyncIO.join();
        }
    }

    // 清理资源
    cleanup();

    if (m_bConnected) {
        // 清除连接标志
        m_bConnected = false;
        COMMONLOG_INFO("Disconnected to server (%s:%u).", m_stSrvConnInfo.strIP.data(), m_stSrvConnInfo.ui16Port);
    } else {
        //std::cerr << ">>>>>> It is not connected to server!" << std::endl;
    }

    return;
}

/*****************************************************
 * 功能：向服务端发送数据
 * 输入参数：
 *      strData：待发送的数据
 * 返回值：
 *      bool：发送结果，true -- 发送成功，false -- 发送失败
 * ***************************************************/
bool TcpClient::sendData(const std::string &strData)
{
    bool bRet = true;
    if (m_bConnected) {
        // 发送数据
        std::lock_guard<std::mutex> lock(m_mt4SendData);

        size_t sTotalSent = 0;
        size_t sDataLen = strData.length();

        while (sTotalSent < sDataLen) {
            int iSent = send(m_fdSock, strData.c_str() + sTotalSent, sDataLen - sTotalSent, 0);
            if (0 <= iSent) {
                sTotalSent += iSent;
            } else {
#ifdef _WIN32
                int error = WSAGetLastError();
                if (error == WSAEWOULDBLOCK) {
                    // 非阻塞模式下缓冲区满，稍等继续发送
                    std::this_thread::sleep_for(std::chrono::milliseconds(10));
                    continue;
                } else {
                    COMMONLOG_WARN("Failed to send data: %d.", error);
                    bRet = false;
                    break;
                }
#else
                if (errno == EWOULDBLOCK || errno == EAGAIN) {
                    // 非阻塞模式下缓冲区满，稍等继续发送
                    std::this_thread::sleep_for(std::chrono::milliseconds(10));
                    continue;
                } else {
                    COMMONLOG_WARN("Failed to send data: %s.", strerror(errno));
                    bRet = false;
                    break;
                }
#endif
            }
        }
    } else {
        COMMONLOG_WARN("%s", "Not connected to server.");
        bRet = false;
    }

    return bRet;
}

/*****************************************************
 * 功能：向服务端读取数据
 * 输出参数：
 *      strData：待读取的数据的缓存区
 * 返回值：
 *      bool：读取结果，true -- 读取成功，false -- 读取失败
 * ***************************************************/
bool TcpClient::recvData(std::string &strData)
{
    bool bRet = false;
    if (m_bConnected) {
        std::lock_guard<std::mutex> lock(m_mt4RecvData);
        fd_set readfds;
        FD_ZERO(&readfds);
        FD_SET(m_fdSock, &readfds);

        struct timeval timeout;
        timeout.tv_sec = 5; // 设置5秒超时
        timeout.tv_usec = 0;

        if (0 < select(m_fdSock + 1, &readfds, nullptr, nullptr, &timeout)) {
            // 有数据可读，可以安全调用recv而不会阻塞
            char szBuffer[RECVDATA_DEFAULT_SIZE] = {0}; // 默认每次最大接收5K数据
            int iReceived = ::recv(m_fdSock, szBuffer, sizeof(szBuffer), 0);
            if (0 < iReceived) {
#if 0
                strData = std::string(szBuffer, iReceived); // 避免数据中间有'\0'被截断
#else
                strData.resize(static_cast<size_t>(iReceived), '\0');
                memcpy((void *)strData.data(), szBuffer, static_cast<size_t>(iReceived));
#endif
                bRet = true;
            }
        }

    } else {
        COMMONLOG_WARN("%s", "Not connected to server.");
    }

    return bRet;
}

/*****************************************************
 * 功能：初始化socket
 * ***************************************************/
bool TcpClient::initSocket()
{
    // 创建socket
#ifdef _WIN32
    m_fdSock = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
#else
    m_fdSock = socket(AF_INET, SOCK_STREAM, 0);
#endif
    if (0 > m_fdSock) {
#ifdef _WIN32
        COMMONLOG_WARN("Failed to create socket: %d.", WSAGetLastError());
#else
        COMMONLOG_WARN("Failed to create socket: %s.", strerror(errno));
#endif
        return false;
    }

    // 设置为非阻塞模式
    //setNonBlocking();

    return true;
}

/*****************************************************
 * 功能：清理资源
 * ***************************************************/
void TcpClient::cleanup()
{
    if (0 <= m_fdSock) {
#ifdef _WIN32
        closesocket(m_fdSock);
#else
        ::close(m_fdSock);
#endif
        m_fdSock = -1;
    }

    return;
}

#if 0
/*****************************************************
 * 功能：设置socket为非阻塞模式
 * ***************************************************/
bool TcpClient::setNonBlocking()
{
#ifdef _WIN32
    u_long mode = 1;
    return (0 == ioctlsocket(m_fdSock, FIONBIO, &mode));
#else
    int flags = fcntl(m_socket, F_GETFL, 0);
    if (flags < 0) { return false; }
    return (fcntl(m_socket, F_SETFL, flags | O_NONBLOCK) == 0);
#endif
}
#endif

/*****************************************************
 * 功能：与服务端进行重新连接
 * 返回值：
 *      bool：连接结果，true -- 连接成功，false -- 连接失败
 * ***************************************************/
bool TcpClient::connectToServer()
{
    // 设置服务器地址
    struct sockaddr_in serverAddr;
    memset(&serverAddr, 0, sizeof(serverAddr));
    serverAddr.sin_family = AF_INET;
    serverAddr.sin_port = htons(m_stSrvConnInfo.ui16Port);

#ifdef _WIN32
    serverAddr.sin_addr.s_addr = inet_addr(m_stSrvConnInfo.strIP.c_str());
    if (serverAddr.sin_addr.s_addr == INADDR_NONE) {
        COMMONLOG_WARN("Invalid IP address: %s.", m_stSrvConnInfo.strIP.data());
        return false;
    }
#else
    if (inet_aton(m_stSrvConnInfo.strIP.c_str(), &serverAddr.sin_addr) == 0) {
        COMMONLOG_WARN("Invalid IP address: %s.", m_stSrvConnInfo.strIP.data());
        return false;
    }
#endif

    // 尝试连接
    if (0 > ::connect(m_fdSock, (struct sockaddr *)&serverAddr, sizeof(serverAddr))) {
#ifdef _WIN32
        int error = WSAGetLastError();
        if (error != WSAEWOULDBLOCK && error != WSAEINPROGRESS) {
            COMMONLOG_WARN("Failed to connect to server, error info: %d.", error);
            return false;
        }
#else
        if (errno != EINPROGRESS) {
            COMMONLOG_WARN("Failed to connect to server, error info: %s.", strerror(errno));
            return false;
        }
#endif
    } else {
        m_bConnected = true; // 设置连接成功的标志
        COMMONLOG_INFO("Connected to server (%s:%u) success.", m_stSrvConnInfo.strIP.data(), m_stSrvConnInfo.ui16Port);
    }

    return true;
}

/*****************************************************
 * 功能：重新连接服务端
 * ***************************************************/
void TcpClient::reconnectServer()
{
    while (m_bRunning) {
        if (!m_bConnected) {
            COMMONLOG_INFO("Attemp reconnecting to server (%s:%u).", m_stSrvConnInfo.strIP.data(), m_stSrvConnInfo.ui16Port);

            // 清除之前的资源
            cleanup();

            // 重新初始化Socket
            initSocket();

            // 连接服务端
            connectToServer();

            // 避免CPU占用过高
            std::this_thread::sleep_for(std::chrono::seconds(3)); // 重连间隔为3s
        } else {
            // 连接成功，跳出循环
            break;
        }
    }

    return;
}

/*****************************************************
 * 功能：异步IO线程业务处理接口
 * ***************************************************/
void TcpClient::onAsyncIOThread()
{
    COMMONLOG_INFO("%s", "Enter in tcp client async IO thread.");

    // 设置select()参数
    struct timeval timeout{};
    timeout.tv_sec = 5; // 设置5秒超时
    timeout.tv_usec = 0;

    fd_set readfds; // 数据读取fds
    while (m_bRunning) {
        if (m_bConnected) {
            FD_ZERO(&readfds);
            FD_SET(m_fdSock, &readfds);

            if (0 > select(m_fdSock + 1, &readfds, nullptr, nullptr, &timeout)) {
                // 异常
                m_bConnected = false;
                if (m_bAutoReconn) {
                    // 执行重连
                    reconnectServer();
                } else {
                    // 不需要重连时，直接跳槽循环
                    COMMONLOG_ERROR("%s", "Connection error, it will exit this async thread!");
                    break;
                }
            } else {
                // 处理读操作，读取数据
                if (FD_ISSET(m_fdSock, &readfds)) {
                    char szBuffer[RECVDATA_DEFAULT_SIZE] = {0}; // 默认每次最大接收5K数据
                    int iRecvLen = ::recv(m_fdSock, szBuffer, sizeof(szBuffer), 0);
                    if (0 < iRecvLen) {
                        std::string strData(szBuffer, static_cast<size_t>(iRecvLen)); // 避免数据中间有'\0'被截断
                        // 调用回调函数，建议通过回调方式获取数据
                        m_mt4PfnCallback.lock();
                        if (m_pfn4RecvDataCb) {
                            m_pfn4RecvDataCb(strData, m_fdSock);
                        }
                        m_mt4PfnCallback.unlock();
                    } else if (0 == iRecvLen) {
                        // 连接已关闭
                        COMMONLOG_WARN("%s", "Connection closed by server!");
                        m_bConnected = false;
                        if (m_bAutoReconn) {
                            // 执行重连
                            reconnectServer();
                        } else {
                            // 不需要重连时，直接跳槽循环
                            break;
                        }
                    } else {
                        m_bConnected = false;
#ifdef _WIN32
                        int error = WSAGetLastError();
                        if (error != WSAEWOULDBLOCK) {
                            COMMONLOG_WARN("Failed to receive data: %d.", error);
                        }
#else
                        if (errno != EWOULDBLOCK && errno != EAGAIN) {
                            COMMONLOG_WARN("Failed to receive data: %s.", strerror(errno));
                        }
#endif
                        if (m_bAutoReconn) {
                            // 执行重连
                            reconnectServer();
                        } else {
                            // 不需要重连时，直接跳槽循环
                            break;
                        }
                    }
                }
            }
        } else {
            if (m_bAutoReconn) {
                // 执行重连
                reconnectServer();
            } else {
                // 无须重连
                break;
            }
        }

        // 短暂休眠以避免过度占用CPU
        std::this_thread::sleep_for(std::chrono::milliseconds(5));
    }

    COMMONLOG_INFO("%s", "Exit to tcp client async IO thread.");
    return;
}

